package com.arthas.client;

import java.io.IOException;

public class ArthasClient {
    /**
     * 连接地址
     */
    private static String host = "127.0.0.1";
    /**
     * 端口
     */
    private static int port = 3658;
    /**
     * 回显提示符开始
     */
    private static String PROMPT_START = "[arthas@";
    /**
     * 回显提示符结束
     */
    private static String PROMPT_END = "]$ ";
    /**
     * 命令返回匹配
     */
    private static String[] PROMPT = new String[]{PROMPT_START, PROMPT_END};
    /**
     * 命令末尾添加纯文本命令和换行
     */
    private static String COMMAND_END = " | plaintext\n";//命令后缀,纯文本返回,不带颜色等格式修饰
    /**
     * 命令执行完回显,会以换行和回车结尾
     */
    private static String COMMAND_RTN_END = " | plaintext\r\n";//命令返回匹配
    /**
     * 客户端
     */
    private TelnetFramework telnet;
    /**
     * 进程id
     */
    private String pid;

    /**
     * 构造客户端并建立连接,默认地址127.0.0.1,端口3658
     */
    public ArthasClient() {
        this(host, port);
    }

    /**
     * 构造客户端并建立连接<br>
     * 需要修改arthas.properties<br>
     * 监听本地某网卡地址,如内网ip<br>
     * arthas.ip=ip<br>
     * 监听所有地址<br>
     * arthas.ip=0.0.0.0<br>
     */
    public ArthasClient(String host, int port) {
        telnet = new TelnetFramework();
        try {
            //建立连接
            telnet.connect(host, port);
            //调整窗口大小,主要是调整宽度稍长一些,避免数据被换行
            telnet.setWindowSize(1024, 500);
            //消耗hello回显,并获取pid
            pid = getPid(telnet.waitFor(PROMPT));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 获取进程id
     *
     * @param rtn hello回显提示
     * @return
     */
    public String getPid(String rtn) {
        //返回信息是匹配PROMPT结尾才返回的
        int start_index = rtn.lastIndexOf(PROMPT_START);
        int end_index = rtn.lastIndexOf(PROMPT_END);
        return rtn.substring(start_index + PROMPT_START.length(), end_index);
    }

    /**
     * 是否连接正常,建议使用前判断
     *
     * @return
     */
    public boolean isConnected() {
        return telnet.isConnected();
    }

    /**
     * 发送命令,处理返回内容
     *
     * @param command 命令
     * @return
     */
    public String sendCommand(String command) {
        String rtn = null;
        if (telnet.isConnected()) {
            try {
                rtn = telnet.write(command + COMMAND_END, PROMPT);
                if (rtn != null) {
                    //回显数据处理,如果命令太长,命令中间可能添加了回车换行,匹配不正确,所以要调整client窗口大小
                    //处理命令回显
                    if (rtn.startsWith(command + COMMAND_RTN_END)) {
                        rtn = rtn.substring((command + COMMAND_RTN_END).length());
                    }
                    //处理回显提示符
                    int index = rtn.lastIndexOf(PROMPT_START + pid + PROMPT_END);
                    if (index != -1) {
                        rtn = rtn.substring(0, index);
                    }
                    //处理尾部换行等空字符
                    rtn = trimEnd(rtn);
                }
            } catch (IOException e) {
                if ("A timeout has occured when trying to read the telnet stream".equals(e.getMessage())) {
                    System.out.println("命令返回超时!");
                } else {
                    e.printStackTrace();
                }
                return null;
            }
        }
        return rtn;
    }

    /**
     * 断开连接
     */
    public void disconnect() {
        if (telnet.isConnected()) {
            try {
                telnet.disconnect();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 处理尾部空字符
     *
     * @param str
     * @return
     */
    private String trimEnd(String str) {
        char[] value = str.toCharArray();
        int len = value.length;
        int st = 0;
        char[] val = value;
        while ((st < len) && (val[len - 1] <= ' ')) {//小于空格的为不可见字符
            len--;
        }
        return ((st > 0) || (len < value.length)) ? new String(val).substring(st, len) : new String(val);
    }

    /**
     * 修改读取超时时间,需要长时间等待操作时进行设置
     *
     * @param time
     */
    public void setTimeOut(long time) {
        telnet.setTimeout(time);
    }


    public static void main(String[] args) throws IOException {
        ArthasClient client = new ArthasClient("192.168.81.128", 3658);
        if (!client.isConnected()) {//本地连接,不存在网络问题
            System.out.println("连接arthas失败,请检查arthas是否启动.");
            return;
        }
        String rtn = client.sendCommand("help");
        System.out.println(rtn);
        rtn = client.sendCommand("classloader");
        System.out.println(rtn);
        rtn = client.sendCommand("mbean");
        System.out.println(rtn);
        client.disconnect();
    }

}
